#!/usr/bin/env python3
"""
流式数据捕获脚本

这个脚本用于捕获一次测试中的HTTP接口返回数据，并保存到日志文件中
"""

import httpx
import json
import asyncio
import time
from datetime import datetime
import os

# API 配置
API_BASE_URL = "http://localhost:8080"
API_ENDPOINT = "/chat"

async def capture_streaming_data():
    """
    捕获流式数据并保存到文件
    """
    print("🎯 开始捕获流式数据...")
    print("=" * 50)
    
    # 测试请求数据
    request_data = {
        "user_id": 1,
        "message": "在平面直角坐标系中，已知点A(0,0)，点B(4,0)，点C(0,3)。\n1. 求三角形ABC的面积\n2. 求三角形ABC的外接圆方程",
        "title": "几何问题测试",
        "description": "用于捕获流式数据格式的测试"
    }
    
    # 创建日志文件名（带时间戳）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = f"streaming_capture_{timestamp}.log"
    
    print(f"📝 请求数据:")
    print(json.dumps(request_data, indent=2, ensure_ascii=False))
    print(f"📁 日志文件: {log_file}")
    print("-" * 50)
    
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            async with client.stream("POST", f"{API_BASE_URL}{API_ENDPOINT}", 
                                    json=request_data) as response:
                
                print(f"📊 响应状态: {response.status_code}")
                print(f"📋 响应头: {dict(response.headers)}")
                print("-" * 50)
                
                if response.status_code != 200:
                    error_text = await response.aread()
                    print(f"❌ 错误响应: {error_text.decode()}")
                    return
                
                # 打开日志文件
                with open(log_file, 'w', encoding='utf-8') as f:
                    # 写入请求信息
                    log_entry = {
                        "timestamp": datetime.now().isoformat(),
                        "type": "REQUEST_INFO",
                        "request_data": request_data,
                        "response_status": response.status_code,
                        "response_headers": dict(response.headers)
                    }
                    f.write(f"=== 请求信息 ===\n")
                    f.write(json.dumps(log_entry, ensure_ascii=False, indent=2))
                    f.write(f"\n\n=== 流式响应数据 ===\n")
                    
                    print("📺 开始接收流式数据:")
                    chunk_count = 0
                    
                    async for chunk in response.aiter_text():
                        if chunk.strip():
                            chunk_count += 1
                            
                            # 记录原始数据块
                            chunk_log = {
                                "chunk_number": chunk_count,
                                "timestamp": datetime.now().isoformat(),
                                "raw_chunk": chunk,
                                "chunk_length": len(chunk)
                            }
                            
                            # 写入日志文件
                            f.write(f"\n--- 数据块 {chunk_count} ---\n")
                            f.write(json.dumps(chunk_log, ensure_ascii=False, indent=2))
                            f.write(f"\n")
                            
                            try:
                                # 尝试解析数据
                                if chunk.startswith("data: "):
                                    data_str = chunk.replace("data: ", "").strip()
                                    if data_str:
                                        data = json.loads(data_str)
                                        
                                        # 记录解析后的数据
                                        parsed_log = {
                                            "chunk_number": chunk_count,
                                            "timestamp": datetime.now().isoformat(),
                                            "parsed_data": data
                                        }
                                        f.write(f"解析后数据:\n")
                                        f.write(json.dumps(parsed_log, ensure_ascii=False, indent=2))
                                        f.write(f"\n")
                                        
                                        # 控制台输出
                                        print(f"  [{chunk_count}] {data.get('type', 'unknown')}: {data.get('content', '')[:100]}{'...' if len(data.get('content', '')) > 100 else ''}")
                                        
                                        if data.get('type') == 'complete':
                                            print(f"✅ 数据捕获完成，共收到 {chunk_count} 个数据块")
                                            break
                                        elif data.get('type') == 'error':
                                            print(f"❌ 处理出错")
                                            break
                                            
                            except json.JSONDecodeError as e:
                                print(f"  [{chunk_count}] 原始数据: {chunk[:100]}{'...' if len(chunk) > 100 else ''}")
                                f.write(f"JSON解析错误: {str(e)}\n")
                            
                            # 刷新文件缓冲区
                            f.flush()
                
                print(f"\n🎉 数据捕获完成！")
                print(f"📁 日志文件已保存: {log_file}")
                print(f"📊 文件大小: {os.path.getsize(log_file)} 字节")
                
                return log_file
                
    except httpx.ConnectError:
        print("❌ 连接错误: 无法连接到服务器")
        print("💡 请确保服务已启动: ./introduction_files/start_conversation_service.sh")
        return None
    except Exception as e:
        print(f"❌ 捕获异常: {str(e)}")
        return None

def analyze_captured_data(log_file: str):
    """
    分析捕获的数据
    """
    if not os.path.exists(log_file):
        print(f"❌ 日志文件不存在: {log_file}")
        return
    
    print(f"\n📊 分析捕获的数据: {log_file}")
    print("=" * 50)
    
    with open(log_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 统计数据块数量
    chunk_count = content.count("--- 数据块")
    print(f"📋 总数据块数: {chunk_count}")
    
    # 统计不同类型的消息
    message_types = []
    lines = content.split('\n')
    for line in lines:
        if '"type":' in line and '"message"' in line:
            message_types.append("message")
        elif '"type":' in line and '"complete"' in line:
            message_types.append("complete")
        elif '"type":' in line and '"error"' in line:
            message_types.append("error")
    
    print(f"📝 消息类型统计:")
    for msg_type in set(message_types):
        count = message_types.count(msg_type)
        print(f"  {msg_type}: {count}")
    
    print(f"📁 文件大小: {os.path.getsize(log_file)} 字节")

async def main():
    """
    主函数
    """
    print("🎯 流式数据捕获工具")
    print("=" * 50)
    
    # 检查服务是否运行
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            # 尝试访问根路径或docs来检查服务
            response = await client.get(f"{API_BASE_URL}/docs")
            if response.status_code == 200:
                print("✅ 服务正在运行")
            else:
                print("⚠️  服务可能未正常运行，但继续尝试...")
    except:
        print("❌ 无法连接到服务")
        print("💡 请先启动服务: python -m uvicorn services.conversation_service_v2:app --host 0.0.0.0 --port 8080")
        return
    
    # 捕获数据
    log_file = await capture_streaming_data()
    
    if log_file:
        # 分析数据
        analyze_captured_data(log_file)
        
        print(f"\n📖 查看完整日志:")
        print(f"cat {log_file}")
        print(f"\n🔍 查看格式化的JSON:")
        print(f"cat {log_file} | grep -A 20 '解析后数据'")

if __name__ == "__main__":
    asyncio.run(main()) 